'''
将pcap文件中的流分离出来，
单独形成文件，
方便后续处理.
'''

'''
外部调用接口：
生成类实例，var= DiversFlow()
调用Divest()方法，传入pcap文件路径和过滤条件（可选）
调用SaveFlow()方法，传入pcap文件路径和保存路径（可选）

结果为：
生成的流文件保存在./processed_dataset/目录下，文件名为pcap文件名_flow_序号.pcap
'''

from scapy.all import *

# 从pcap文件中提取流的封装类
class DivestFlow:
    def __init__(self) -> None:
        self.id = []             # 用于存储流的五元组id
        self.flow_data = []      # 用于存储流的数据包
        self.flow_count = 0      # 用于记录流的数量
        
    # 获取数据包的五元组
    def Get5tuple(self, pkt):
        return (pkt["IP"].proto, pkt["IP"].src, pkt["IP"].dst, pkt.payload.sport, pkt.payload.dport)
    
    # 处理数据包
    def Divest(self, pcap_file, custom_filter = ""):
        pkts = sniff(offline=pcap_file, filter = custom_filter)
        
        for pkt in pkts:
            # 获取数据包的五元组
            five_tuple = self.Get5tuple(pkt)
            five_tuple_reverse = (five_tuple[0], five_tuple[2], five_tuple[1], five_tuple[4], five_tuple[3])
            
            # 如果该五元组没有出现过，则创建新的流
            if five_tuple in self.id:
                index = self.id.index(five_tuple)
                self.flow_data[index].append(pkt)
            elif five_tuple_reverse in self.id:
                index = self.id.index(five_tuple_reverse)
                self.flow_data[index].append(pkt)
                
            else:
                self.id.append(five_tuple)
                self.flow_data.append([pkt])
                self.flow_count += 1
                
    def SaveFlow(self, pcap_file, save_path = "./processed_dataset/"):
        for i in range(self.flow_count):
            flow_file = save_path + pcap_file.split("/")[-1].split(".")[0] + "_flow_" + str(i) + ".pcap"
            # 将包内容逐个写入
            for pkt in self.flow_data[i]:
                wrpcap(flow_file, pkt, append=True)
    

if __name__ == "__main__":
    divest_flow = DivestFlow()
    pcap_file = "./dataset/browsing.pcap"
    
    divest_flow.Divest(pcap_file, "tcp")
    '''
    divest_flow.SaveFlow(pcap_file)
    print("Total %d flows saved." % divest_flow.flow_count)
    '''        
    